from qutip import *
import numpy as np
import matplotlib.pyplot as plt
sys.path.append('D:\Dropbox\Dokumente\PI3\Code\Python\quantum-tools')
import qutip_enhanced
bloch = Bloch3d()
qte = qutip_enhanced.QutipEnhanced()

def h(nu_rf, hf, omega_e=0, phase_e=0, omega_n=0, phase_n=0):
    a = 2*np.pi*nu_rf/2.
    b = 2*np.pi*(nu_rf - hf)/2.
    c = 0.5*2*np.pi*omega_e*np.exp(1j*phase_e)
    d = 0.5*2*np.pi*omega_n*np.exp(1j*phase_n)
    h = Qobj([[a, 0,  c.conj(), 0],
             [0,  -a, 0,        c.conj()],
             [c,  0,  b,        d.conj()],
             [0,  c,  d,        -b]])
    h.dims = [[2, 2], [2, 2]]
    return h

def rho0():
    rho0_e = ket2dm(basis(2, 0))
    rho0_n = ket2dm(basis(2, 0))
    return tensor(rho0_e, rho0_n)

gamma = 1/5850.
collapse_oper_t1 = tensor(np.sqrt(gamma)* (jmat(0.5, 'x') + jmat(0.5, 'y')), qeye(2)) #found by trying.. it just gives the right decay
gamma_t2 = 1/2500
collapse_oper_t2 = tensor(np.sqrt(gamma_t2)*jmat(0.5, 'z'), qeye(2))

def evolve(rho, length_mus, **kwargs):
    rho = mesolve(h(**kwargs), rho, np.linspace(0, length_mus, 2), [], []).states[-1]
    return rho

def KDD5(rho, shift, tau, omega_e, omega_n, nkdd5=0, **kwargs):
    #tau
    rho = evolve(rho, tau, **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 6. + shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, **kwargs)
    rho = evolve(rho, tau, **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, **kwargs)
    rho = evolve(rho, tau, **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 2. + shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, **kwargs)
    rho = evolve(rho, tau, **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=shift, **kwargs)
    #two tau
    rho = evolve(rho, tau, **kwargs)
    rho = evolve(rho, tau, **kwargs)
    #e-pi
    rho = evolve(rho, 0.5/omega_e, omega_e=omega_e, phase_e=np.pi / 6. + shift, **kwargs)
    #tau
    rho = evolve(rho, tau, **kwargs)
    return rho

def sequence(**kwargs):
    rho = rho0()
    rho = evolve(rho, length_mus=.25/kwargs['omega_e'], omega_e=kwargs['omega_e'], hf=kwargs['hf'], nu_rf=kwargs['nu_rf'])
    rho = KDD5(rho, shift=0.0, nkdd5=0, **kwargs)
    rho = KDD5(rho, shift=np.pi / 2., nkdd5=1, **kwargs)
    rho = KDD5(rho, shift=0.0, nkdd5=2, **kwargs)
    rho = KDD5(rho, shift=np.pi / 2., nkdd5=3, **kwargs)
    rho = evolve(rho, length_mus=.25/kwargs['omega_e'], omega_e=kwargs['omega_e'], hf=kwargs['hf'], nu_rf=kwargs['nu_rf'])
    return rho


_omega_e_ = 1/0.2
_omega_n_ = 0
_hf_ = 0
_nu_rf_ = 0

exp_z = []
tau_l = np.linspace(0, 500, 50)[0:1]
for _tau_ in tau_l:
    kwargs = dict(tau=_tau_, omega_e=_omega_e_, omega_n=_omega_n_, hf=_hf_, nu_rf=_nu_rf_)
    exp_z.append(expect(1/2.*sigmaz(), sequence(**kwargs).ptrace(0)))

import matplotlib.pyplot as plt
plt.plot(tau_l, exp_z)
plt.show()